#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Fri Mar  2 16:09:13 2018

@author: cc
"""
import numpy as np
from skimage.io import imsave
import matplotlib.pyplot as plt
from sklearn.metrics import confusion_matrix
from utils import get_data
from utils import add_white_space
from utils import patch_to_img


import keras
from keras.models import Sequential
from keras.layers import Dense, Dropout, Flatten
from keras.layers import Conv2D, MaxPooling2D


def build_cnn(img_cols, img_rows, nb_class=3, nb_channel=1):
    input_shape = (img_cols, img_rows, nb_channel)
    model = Sequential()
    model.add(Conv2D(32, kernel_size=(3, 3),
                     activation='relu',
                     input_shape=input_shape))
    model.add(Conv2D(64, (3, 3), activation='relu'))
    model.add(MaxPooling2D(pool_size=(2, 2)))
    model.add(Dropout(0.25))
    model.add(Flatten())
    model.add(Dense(128, activation='relu'))
    model.add(Dropout(0.5))
    model.add(Dense(nb_class, activation='softmax'))

    model.compile(loss="categorical_crossentropy",
                  optimizer="Adam",
                  metrics=['accuracy'])
    return model


def show_val_samples(target='correct'):
    labels = ['sitting', 'standing', 'lying']
    disp_rows = len(labels)
    disp_cols = 8

    test_label = np.argmax(y_test, axis=1)
    pred_test = model.predict_classes(X_test)
    if target == 'correct':
        pred_y = test_label[pred_test == test_label]
        pred_X = X_data[pred_test == test_label]
    else:
        pred_y = test_label[pred_test != test_label]
        pred_X = X_data[pred_test != test_label]

    results = []
    for i in range(0, disp_rows):
        for j in range(0, disp_cols):
            sample = pred_X[pred_y == i]
            print(i,j,sample.shape)
            sample = sample[j, :, :, 0]
            sample = (sample*255).astype('uint8')
            results.append(add_white_space(sample, 5))
    img = patch_to_img(results, disp_rows, disp_cols)
    imsave('result/{}_sample.png'.format(target), img)
    return 0

def show_pred_distribution():
    test_label = np.argmax(y_test, axis=1)
    pred_test = model.predict_classes(X_test)
    pred_correct = test_label[pred_test == test_label]
    pred_wrong = test_label[pred_test != test_label]
    bins = [0, 1, 2, 3]
    plt.figure()
    plt.hist(test_label, bins=bins, label='all')
    plt.hist(pred_correct, bins=bins, label='correct')
    plt.hist(pred_wrong, bins=bins, label='wrong')
    plt.legend()
    plt.ylabel('Number of samples')
    plt.xlabel(['sitting','standing','lying'])
    plt.title('Prediction distribution')
    plt.savefig('result/test_prediction_.png')

if __name__ == "__main__":
    # load data
    img_cols, img_rows = 56, 56
    nb_class, nb_channel, ratio = 3, 1, 0.7
    X_data, y_data = get_data('trainval', img_cols, img_rows)
    y_data = keras.utils.to_categorical(y_data, nb_class)

    X_test, y_test = get_data('test', img_cols, img_rows)
    y_test = keras.utils.to_categorical(y_test, nb_class)

    model = build_cnn(img_cols, img_rows, nb_class, nb_channel)
    model.fit(X_data, y_data,
              batch_size=32,
              epochs=10,
              verbose=1,
              validation_data=(X_test, y_test))

    score_train = model.evaluate(X_data, y_data)
    score_test = model.evaluate(X_test, y_test)

    print('Train loss:', score_train[0])
    print('Train accuracy:', score_train[1])
    print('Test loss:', score_test[0])
    print('Test accuracy:', score_test[1])

    for target in ['correct', 'wrong']:
        show_val_samples(target)
    show_pred_distribution()

    pred_data = model.predict_classes(X_data)
    pred_test = model.predict_classes(X_test)
    y_data = np.argmax(y_data, axis=1)
    y_test = np.argmax(y_test, axis=1)
    matrix_data = confusion_matrix(y_data, pred_data)
    matrix_test = confusion_matrix(y_test, pred_test)

    import seaborn as sns
    plt.figure(figsize = (10,7))
    sns.set(font_scale=1.4)
    sns.heatmap(matrix_data, annot=True, annot_kws={"size": 12})
    plt.xlabel("Predit Labels", fontsize=16)
    plt.ylabel("True Labels", fontsize=16)
    plt.title("Sitting(0);Standing(1);Lying(2)")
    plt.savefig('result/train_confusion.png')

    plt.figure(figsize = (10,7))
    sns.set(font_scale=1.4)
    sns.heatmap(matrix_test, annot=True, annot_kws={"size": 12})
    plt.xlabel("Predit Labels", fontsize=16)
    plt.ylabel("True Labels", fontsize=16)
    plt.title("Sitting(0);Standing(1);Lying(2)")
    plt.savefig('result/test_confusion.png')